# app/routes/prediction_routes.py
from flask import Blueprint, request, jsonify, send_from_directory, session, current_app
import uuid
import os
import cv2
import numpy as np
from app.services.prediction_service import process_image, generate_report
from app.models import PredictionHistory, PredictionReport, User
from app import db
import json
from app.utils import login_required
import logging

prediction_bp = Blueprint('prediction', __name__)

@prediction_bp.route('/predict', methods=['POST'])
@login_required
def predict():
    if 'files' not in request.files:
        return jsonify({'error': 'No files uploaded'}), 400
    files = request.files.getlist('files')
    if not files:
        return jsonify({'error': 'No files selected'}), 400
    input_urlList = []
    mask_urlList = []
    os.makedirs('uploads', exist_ok=True)
    os.makedirs('static', exist_ok=True)
    from app.model_loader import preloaded_models
    if preloaded_models is None or len(preloaded_models) == 0:
        for file in files:
            if file.filename == '':
                continue
            upload_uuid = uuid.uuid4()
            upload_path = f'uploads/input_image_{upload_uuid}.jpg'
    for index, file in enumerate(files):
        if file.filename == '':
            continue
        # 使用 UUID 生成唯一的上传文件名
        upload_uuid = uuid.uuid4()
        upload_path = f'uploads/input_image_{upload_uuid}.jpg'
        file.save(upload_path)
        try:
            # 使用预加载的模型进行处理
            mask = process_image(upload_path)  # 不再传递 preloaded_models 参数
            if mask is None:
                continue
            # 使用 UUID 生成唯一的掩码图像文件名
            mask_uuid = uuid.uuid4()
            mask_path = f'static/mask_image_{mask_uuid}.png'
            mask_colored = cv2.applyColorMap((mask * 50).astype(np.uint8), cv2.COLORMAP_JET)
            cv2.imwrite(mask_path, mask_colored)
            input_url = f'/uploads/input_image_{upload_uuid}.jpg'
            mask_url = f'/static/mask_image_{mask_uuid}.png'
            input_urlList.append(input_url)
            mask_urlList.append(mask_url)
        except Exception as e:
            logging.error(f"Error processing image: {str(e)}")

    passed_qty = len(input_urlList)
    sum_qty = len(files)
    passed_percent = passed_qty / sum_qty if sum_qty > 0 else 0

    input_urls_json = json.dumps(input_urlList)
    mask_urls_json = json.dumps(mask_urlList)

    report_content, pdf_path, _, _ = generate_report(input_urls_json, mask_urls_json, passed_percent, len(files))
    report = PredictionReport(report_content=report_content, passed_percent=passed_percent,
                              input_urls=input_urls_json, mask_urls=mask_urls_json)
    db.session.add(report)
    db.session.commit()
    user = User.query.get(session['user_id'])
    history = PredictionHistory(user=user, report=report)
    db.session.add(history)
    db.session.commit()

    return jsonify({
        '报告': json.loads(report_content),
        'history_id': history.id,
        'pdf_url': f'/static/prediction_report_{history.id}.pdf'
    })


@prediction_bp.route('/uploads/<filename>')
def uploaded_file(filename):
    return send_from_directory('uploads', filename)

@prediction_bp.route('/history', methods=['GET'])
@login_required
def view_history():
    if 'user_id' not in session:
        return jsonify({'error': '请先登录'}), 401
    user = User.query.get(session['user_id'])
    histories = user.prediction_histories.all()
    history_list = []
    for history in histories:
        history_list.append({
            'history_id': history.id,  # 添加 history_id 字段
            'user_id': history.user_id,
            'prediction_time': history.prediction_time.strftime('%Y-%m-%d %H:%M:%S'),
            'report': history.report.get_report_dict()
        })
    return jsonify(history_list)

@prediction_bp.route('/download_report/<int:history_id>', methods=['GET'])
@login_required
def download_report(history_id):
    if 'user_id' not in session:
        return jsonify({'error': '请先登录'}), 401
    history = PredictionHistory.query.get(history_id)
    if not history or history.user_id != session['user_id']:
        return jsonify({'error': '未找到该历史报告或无权限访问'}), 404
    report = history.report
    input_urls = report.input_urls
    mask_urls = report.mask_urls
    passed_percent = report.passed_percent

    # 重新生成报告
    _, pdf_path, _, _ = generate_report(input_urls, mask_urls, passed_percent, len(json.loads(input_urls)))

    # 检查文件是否存在
    if os.path.exists(pdf_path):
        logging.info(f"PDF report found: {pdf_path}")
        # 获取文件所在目录和文件名
        directory = os.path.dirname(pdf_path)
        filename = os.path.basename(pdf_path)
        return send_from_directory(directory, filename, as_attachment=True)
    else:
        logging.error(f"PDF report not found: {pdf_path}")
        return jsonify({'error': '报告文件未找到，请稍后重试'}), 404